package com.demo.config.apache.NLP;

import opennlp.tools.namefind.NameFinderME;
import opennlp.tools.namefind.TokenNameFinderModel;
import opennlp.tools.util.Span;

import java.io.FileInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 暂时是废的
 */
public class NameFinder {

	// 初始化命名实体识别器
	private NameFinderME personFinder;

	private NameFinderME locationFinder;

	public NameFinder() throws IOException {
		// 加载人物姓名识别模型
		FileInputStream personModelIn = new FileInputStream("D:\\file\\en-ner-person.bin");
		TokenNameFinderModel personModel = new TokenNameFinderModel(personModelIn);
		personFinder = new NameFinderME(personModel);

		// 加载地理位置识别模型
		FileInputStream locationModelIn = new FileInputStream("D:\\file\\en-ner-location.bin");
		TokenNameFinderModel locationModel = new TokenNameFinderModel(locationModelIn);
		locationFinder = new NameFinderME(locationModel);
	}

	public List<String> findEntities(String text) {
		// 识别人物实体
		Span[] personSpans = personFinder.find(text.split(" "));

		// 识别地理位置实体
		Span[] locationSpans = locationFinder.find(text.split(" "));

		// 创建所有实体的列表
		List<String> entities = new ArrayList<>();
		for (Span span : personSpans) {
			entities.add(text.substring(span.getStart(), span.getEnd()));
		}
		for (Span span : locationSpans) {
			entities.add(text.substring(span.getStart(), span.getEnd()));
		}

		return entities;
	}

	public static void main(String[] args) throws IOException {
		String text = "John Smith 访问了 London 和 Paris 上周。";

		NameFinder nameFinder = new NameFinder();
		List<String> entities = nameFinder.findEntities(text);

		System.out.println("实体：" + entities);
	}

}